import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import argparse
import logging
import os
import joblib
from dataset import DataProcessor  
from my_logging import setup_logging
import logging
from model import LSTMModel

# 参数解析
def parse_args():
    parser = argparse.ArgumentParser(description='LSTM Time Series Testing')
    
    # 数据参数
    parser.add_argument('--data_path', type=str, default='data/airline-passengers.csv',
                       help='Path to the dataset file')
    
    # 模型参数
    parser.add_argument('--input_size', type=int, default=1,
                       help='Size of input features')
    parser.add_argument('--hidden_size', type=int, default=50,
                       help='Number of hidden units in LSTM')
    parser.add_argument('--output_size', type=int, default=1,
                       help='Size of output prediction')
    parser.add_argument('--num_layers', type=int, default=1,
                       help='Number of LSTM layers')
    parser.add_argument('--seq_length', type=int, default=12,
                       help='Sequence length for time series')
    
    # 系统参数
    parser.add_argument('--model_dir', type=str, default='models',
                       help='Directory containing saved models')
    parser.add_argument('--log_dir', type=str, default='logs',
                       help='Directory to save logs')
    
    return parser.parse_args()


# 测试流程
def test_main():
    args = parse_args()
    logger = setup_logging(args.log_dir)
    logger.info("Starting LSTM Time Series Testing")
    
    # 加载模型和scaler
    model_path = os.path.join(args.model_dir, 'lstm_model.pth')
    scaler_path = os.path.join(args.model_dir, 'scaler.pkl')
    
    if not os.path.exists(model_path) or not os.path.exists(scaler_path):
        logger.error("Model or scaler file not found!")
        return
    
    # 初始化模型
    model = LSTMModel(args)
    model.load_state_dict(torch.load(model_path))
    model.eval()
    
    # 加载scaler
    scaler = joblib.load(scaler_path)
    
    # 加载数据
    df = pd.read_csv(args.data_path, parse_dates=['Month'], index_col='Month')
    values = df['Passengers'].values.astype('float32')
    
    # 准备测试数据（使用训练时的分割比例）
    split_index = int(len(values) * 0.8)  # 与训练时相同的比例
    test = values[split_index:]
    
    # 标准化测试数据
    test_scaled = scaler.transform(test.reshape(-1, 1)).flatten()
    
    # 创建测试序列
    X_test, y_test = [], []
    for i in range(len(test_scaled)-args.seq_length):
        X_test.append(test_scaled[i:i+args.seq_length])
        y_test.append(test_scaled[i+args.seq_length])
    X_test = np.array(X_test)
    y_test = np.array(y_test)
    
    # 转换为Tensor
    X_test = torch.FloatTensor(X_test).unsqueeze(-1)
    y_test = torch.FloatTensor(y_test).unsqueeze(-1)
    
    # 测试集评估
    with torch.no_grad():
        test_pred = model(X_test)
        test_pred_orig = scaler.inverse_transform(test_pred.numpy())
        test_true_orig = scaler.inverse_transform(y_test.numpy())
        
        mse = np.mean((test_pred_orig - test_true_orig)**2)
        rmse = np.sqrt(mse)
        mae = np.mean(np.abs(test_pred_orig - test_true_orig))
    
    logger.info("\nTest Results:")
    logger.info(f"Test RMSE: {rmse:.2f} passengers")
    logger.info(f"Test MAE: {mae:.2f} passengers")
    
    # 可视化预测结果
    plt.figure(figsize=(12, 6))
    plt.plot(test_true_orig, label='True Passengers')
    plt.plot(test_pred_orig, label='Predicted Passengers')
    plt.title('Airline Passengers Prediction on Test Set')
    plt.xlabel('Time Steps')
    plt.ylabel('Passengers')
    plt.legend()
    plt.grid(True)
    plt.savefig('test_predictions.png')
    plt.close()
    
    logger.info("Testing completed. Results saved to test_predictions.png")

if __name__ == '__main__':
    test_main()